"""
连接mysql示例
pip install mysqlclient
"""
import re
from datetime import datetime
from operator import itemgetter
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDatabaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tools import Tool
from models import get_ds_model_client

#数据库配置
HOSTNAME ='192.168.46.128'
PORT ='3306'
DATABASE = 'test'
USERNAME = 'root'
PASSWORD ='root'
MYSQL_URI ='mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME,PASSWORD,HOSTNAME,PORT,DATABASE)

db = SQLDatabase.from_uri(MYSQL_URI)

def get_table_name():
    return db.get_usable_table_names()

def get_current_time():
    curr_time = datetime.now()
    format_time = curr_time.strftime("%Y-%m-%d %H:%M:%S")
    return f"当前时间:{format_time}"

get_table_names_tool = Tool(
    name="get_table_names",
    func=get_table_name,
    description="获取数据库中所有表名"
)

get_curr_time_tool = Tool(
    name="get_current_time",
    func=get_current_time,
    description="获取当前时间"
)

client = get_ds_model_client()
client_with_tools = client.bind_tools([get_table_names_tool, get_curr_time_tool])

# resp = client_with_tools.invoke([HumanMessage(content="请查询字典表中所有数据")])
# print(resp)

class SQLCleaner(StrOutputParser):
    def parse(self, text: str) -> str:
        pattern = r'```sql(.*?)```'
        match = re.search(pattern, text, re.DOTALL)
        if match:
            sql = match.group(1).strip()
            # 某些大模型还会产生类似'SQLQuery:'前缀，必须去除
            sql = re.sub(r'^SQLQuery:', '', sql).strip()
            return sql
        # 某些大模型还会产生类似'SQLQuery:'前缀，必须去除
        text = re.sub(r'^SQLQuery:', '', text).strip()
        return text

sql_chain = create_sql_query_chain(client, db) | SQLCleaner()
resp = sql_chain.invoke(input={"question": "请查询字典表中所有数据"})
# print(resp)
# print("**"*15)

answer_prompt = PromptTemplate.from_template(
    """给定以下用户问题、可能的SQL语句和SQL执行后的结果，回答用户问题
    Question: {question}
    SQL Query: {query}
    SQL Result:{result}
    回答:"""
)

execute_sql_tool = QuerySQLDatabaseTool(db=db)

chain = (RunnablePassthrough.assign(query=sql_chain)
         .assign(result=itemgetter("query") | execute_sql_tool) | answer_prompt | client | StrOutputParser())
result = chain.invoke(input={"question": "请问字典表中有几条数据"})
print(result)
